from __future__ import with_statement

import sys
import os
import time
import shutil
import unittest
import traceback
import server
from kineta import *


class JoinTests(unittest.TestCase):

    def setUp(self):
        self._dbdir = "db"
        if os.path.isdir(self._dbdir):
            shutil.rmtree(self._dbdir)
        self.open()
        self.db.insert_bulk("w", ["a", "b"],
                                 [[1, 2],
                                  [3, 4],
                                  [5, 6],
                                  [7, 8]])
        self.db.insert_bulk("x", ["a", "c"],
                                 [[1, 9],
                                  [1, 10],
                                  [5, 12],
                                  [7, 12]])
        self.db.insert_bulk("y", ["c", "d"],
                                 [[9, 6],
                                  [9, 7],
                                  [9, 11],
                                  [10, 12]])
        self.db.insert_bulk("z", ["zb", "e"],
                                 [[2, 1],
                                  [2, 2],
                                  [10, 3],
                                  [10, 4]])

    def _open(self):
        db = FileDB(self._dbdir)
        db.set_bg_error_cb(self.bgerror)
        return db

    def open(self):
        self.db = self._open()

    def bgerror(self, exc_info):
        traceback.print_exception(*exc_info)

    def testBasicJoin(self):
        with self.db.select(table=["w", "wa"], fields=["a", "b", ["xa"]],
                            joins={
                                "table": ["x", "xa"],
                                "fields": ["c"],
                                "on": {"a": "a"}}) as reader:
            res = list(reader)
            names = reader.field_names
        self.assertEqual(len(res), 4)
        self.assert_([1, 2, 9] in res)
        self.assert_([1, 2, 10] in res)
        self.assert_([5, 6, 12] in res)
        self.assert_([7, 8, 12] in res)
        self.assertEqual(names, [["wa", "a"], ["wa", "b"], ["xa", "c"]])

    def testMergeJoin(self):
        self.db.create_index("w", [('a', 'a')])
        self.db.create_index("x", [('a', 'a')])
        self.db.wait_for_bg_tasks()
        q = dict(table=["w", "wa"], fields=["a", "b", ["xa"]],
                            joins={
                                "table": ["x", "xa"],
                                "fields": ["c"],
                                "on": {"a": "a"}})
        plan = self.db.explain(**q)
        self.assertEqual(plan['strategy'], 'merge_join')
        with self.db.select(**q) as reader:
            res = list(reader)
            names = reader.field_names
        self.assertEqual(len(res), 4)
        self.assert_([1, 2, 9] in res)
        self.assert_([1, 2, 10] in res)
        self.assert_([5, 6, 12] in res)
        self.assert_([7, 8, 12] in res)
        self.assertEqual(names, [["wa", "a"], ["wa", "b"], ["xa", "c"]])

    def testNoFields(self):
        with self.db.select(table="w", fields=["a", ["fake"], "b", ["xa"]],
                            joins={
                                "table": ["x", "xa"],
                                "on": {"a": "a"}}) as reader:
            res = list(reader)
            names = reader.field_names
        self.assertEqual(len(res), 4)
        self.assert_([1, 2] in res)
        self.assert_([1, 2] in res)
        self.assert_([5, 6] in res)
        self.assert_([7, 8] in res)
        self.assertEqual(names, [["w", "a"], ["w", "b"]])

    def testOuterNestedJoin(self):
        query = dict(table=["w"], fields=["a", "b", ["xa"]],
                     joins=[{"table": ["x", "xa"],
                             "fields": ["c", ["y"]],
                             "on": {"a": "a"},
                             "joins": {"table": "y",
                                       "fields": ["c", "d"],
                                       "on": {"c": "c"},
                                       "where": {"d": ('<', 10)},
                                       "outer": True,
                                       "order": "d"},
                             "order": "c"},
                            {"table": "z",
                             "on": {"b": "zb"}}])
        with self.db.select(**query) as reader:
            res = list(reader)
            names = reader.field_names
        self.assertONJ(res, names)

        self.db.create_index("x", ["c"])
        self.db.create_index("y", ["c"])
        self.db.wait_for_bg_tasks()

        plan = self.db.explain(**query)
        self.assertEqual(plan['sub_plan']['left']['right']['strategy'],
                         'merge_join')

        with self.db.select(**query) as reader:
            res = list(reader)
            names = reader.field_names
        self.assertONJ(res, names)

    def assertONJ(self, res, names):
        self.assertEqual(len(res), 6)
        self.assertEqual(res[0], [1, 2, 9, 9, 6])
        self.assertEqual(res[1], [1, 2, 9, 9, 6])
        self.assertEqual(res[2], [1, 2, 9, 9, 7])
        self.assertEqual(res[3], [1, 2, 9, 9, 7])
        self.assertEqual(res[4], [1, 2, 10, None, None])
        self.assertEqual(res[5], [1, 2, 10, None, None])
        self.assertEqual(names, [["w", "a"], ["w", "b"], ["xa", "c"],
                                 ["y", "c"], ["y", "d"]])

    def testOuterJoin(self):
        with self.db.select(table="x", fields=["a", "c", ["y"]],
                            joins={"table": "y",
                                   "fields": ["c", "d"],
                                   "on": {"c": "c"},
                                   "where": {"d": ('<', 10)},
                                   "outer": True,
                                   "order": "d"},
                            order=["a", "c"]) as reader:
            res = list(reader)
            names = reader.field_names
        self.assertEqual(len(res), 5)
        self.assertEqual(res[0], [1, 9, 9, 6])
        self.assertEqual(res[1], [1, 9, 9, 7])
        self.assertEqual(res[2], [1, 10, None, None])
        self.assertEqual(res[3], [5, 12, None, None])
        self.assertEqual(res[4], [7, 12, None, None])
        self.assertEqual(names, [["x", "a"], ["x", "c"],
                                 ["y", "c"], ["y", "d"]])

    def tearDown(self):
        self.db.close()


class JoinTestsRemote(JoinTests):

    def _open(self):
        return RemoteDB("http://localhost:20000/db")

    def open(self):
        self.server = server.TestServer()
        self.server.app.databases['db'].set_bg_error_cb(self.bgerror)
        self.db = self._open()

    def tearDown(self):
        self.server.quit()
        try:
            self.db.list_tables()
        except:
            pass
        self.db.close()
        self.server.join()
